Write a custom CUDA kernel to optimize `torch.nn.HuberLoss`.


The original operation is defined by the formula:
loss(x, y)_i =
  - 0.5 * (x_i - y_i)^2, if |x_i - y_i| < delta
  - delta * (|x_i - y_i| - 0.5 * delta), otherwise

This is followed by a reduction operation over all elements ('none', 'mean', or 'sum').

**Problem Analysis:**
The standard PyTorch implementation of HuberLoss is memory-bound. It executes a chain of element-wise operations (subtraction, absolute value, comparison, multiplication, etc.) and a final reduction. Each step materializes a full-sized intermediate tensor in global memory, which is immediately read back by the next operation. This results in excessive memory traffic and multiple kernel launch overheads, which are the primary performance bottlenecks.

**Optimization Strategy: Fused Computation and Parallel Reduction**

The goal is to create a CUDA implementation that fuses all stages into one or two kernel launches.

1.  **Fusion for `reduction='none'`**:
    A single element-wise kernel is implemented. Each thread is assigned to one element of the input tensors. It performs the entire Huber Loss calculation (diff, abs, condition, formula) in registers and writes the final result directly to the output tensor. This completely eliminates intermediate memory traffic.

2.  **Fusion for `reduction='mean'` or `'sum'`**:
    A highly-optimized, two-stage parallel reduction strategy is employed:
    *   **Kernel 1 (Calculation & Block-Level Reduction)**: This kernel is launched with a grid size large enough to cover all elements.
        - Each thread computes the Huber loss for one or more elements.
        - The results within a thread block are then efficiently summed up using **shared memory** in a tree-like reduction pattern. This avoids slow global memory atomics.
        - The first thread of each block writes its block's partial sum to a temporary intermediate buffer in global memory.
    *   **Kernel 2 (Final Reduction)**: A second, much smaller kernel (often a single block) is launched. It reads the partial sums from the intermediate buffer and performs the final reduction, again using shared memory, to produce a single scalar result.
    *   For `reduction='mean'`, the final sum is divided by the total number of elements.

This comprehensive fusion strategy minimizes global memory access to a single pass over the input data, drastically reducing bandwidth usage and kernel launch overhead, leading to significant performance gains.

You are given the following architecture:
import torch
import torch.nn as nn

# --- 用于基准测试的配置 ---
batch_size = 512
dim = 4096
DELTA = 1.0
REDUCTION = 'mean'

class Model(nn.Module):
    """
    使用 PyTorch 内置的 torch.nn.HuberLoss 作为基准模型。
    """
    def __init__(self, delta=1.0, reduction='mean'):
        super(Model, self).__init__()
        self.loss_fn = nn.HuberLoss(delta=delta, reduction=reduction)
    
    def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
        return self.loss_fn(input_tensor, target_tensor)

def get_inputs():
    """
    生成用于测试的输入张量。
    """
    input_tensor = torch.randn(batch_size, dim, dtype=torch.float32)
    target_tensor = input_tensor + torch.randn(batch_size, dim, dtype=torch.float32) * 0.5
    return [input_tensor.contiguous(), target_tensor.contiguous()]

def get_init_inputs():
    """
    提供模型初始化所需的参数。
    """
    return [DELTA, REDUCTION]